# main(ion, nsp, fwhm, s2n, minsig, dv, safe, outprefix, specdir,
#         infilename, chsqstop, debug=1):
#
# input ion         = ('c4','h1','o6','si4','n5')
#       nsp         = None             --  fit all spectra,
#                   = integer n        -- fit first n spectra
#                   = [n1,n2,n3,....]  -- fit spectra n1,n2,n3 (first spec is 1)
#       fwhm        = float: intrumental broadening fwhm in km/s
#       s2n         = float: SNR per rebinned pixel at continuum
#       minsig      = float: sigma of gaussian noise component (rest is poisson)
#       dv          = float: velocity width of mock spectra pixels
#       safe        = float: safe distance in Mpc/h from centre to edge
#       outprefix   = string: prefix for spectrum filenames, e.g. 'runA_'
#       specdir     = string: name of output directory
#       infilename  = string: name of hdf5 file containing gimic spectra
#       chsqstop    = float: stopping vpfit once normalised reaches this
#
#
# output: outprefix_ion_spectrum_number.txt contains spectrum
#         outprefix_ion_spectrum_number.f26 contains vpfitted model
#
#
# given a spectrum, runs vpfit to recover line parameters. For now
# assume we have a continuum.
#
# usage: run -i run.py ex is like @ in idl: "-i" means include the
# interpreter's namespace when running the script
#
#
# examples: 
#     main('c4', 1, 6.6, 50, 0.01, 3, 11.0, 'runA_', 'HiRes/Spectra/Random/Sigma0/snap_020/', '/data/rw15/tt/GIMIC/HiRes/Spectra/Random/Sigma0/snap_020/Spectrum.Part_los_z3.060.hdf5',1.2)
#     main('h1', 1, 6.6, 50, 0.01, 3, 11.0, 'runA_', 'HiRes/Spectra/Random/Sigma0/snap_020/', '/data/rw15/tt/GIMIC/HiRes/Spectra/Random/Sigma0/snap_020/Spectrum.Part_los_z3.060.hdf5',1.2)
#
# plot fit using rdgen (hint: provide atom.dat with fitted transition parameters)
# for example:
#%rdgen
#    rd Sigma0c4_001.txt
#    gp 
#        Sigma0c4_001.26
#    6.6 km/s   (FWHM)
#    pg         (to plot)

import numpy as np
import os
import subprocess
import sys
import tables
import astro.utilities

################## Spectrum and fitting parameters #####################

ion = 'h1'
# numbers of the spectra to read from file. if None, then read all the
# spectra. If a single integer, read the first nsp spectra.
nsp = None
fwhm = 6.6            # instrumental broadening fwhm in km/s
s2n = 50              # desired signal to noise ratio at continuum
minsig = 0.01         # snr at zero flux. 
dv = 3.0              # rebin to pixels with this velocity width in km/s
safe = 11.0           # safe region radius in h**-1 Mpc
outprefix = 'runA_'
# specdir = 'LoRes/Spectra/QSO-correlations/Sigma0/snapshot_046'
# infilename = ('/media/cosma/data/rw15/nhmc/GIMIC/LoRes/Spectra/'
#               'QSO-correlations/Sigma0/snapshot_046/'
#               'Spectrum.Part_los_z0.250.hdf5')
specdir = 'LoRes/Spectra/QSO-correlations/Sigmam2/snapshot_027'
infilename = os.path.join('/media/cosma/data/rw15/nhmc/GIMIC',specdir,
                          'Spectrum.Part_los_z1.000.hdf5')
chsqstop = 1.1           # stopping criterion for vpfit

################## End Spectrum and fitting parameters #####################

########### Info needed for run_vpfit section ##############

vpsetup_txt = """\
lastchtied v
! ---- parameter limits ----
bvalmax 50000.0          Maximum b value (can have two - metals, HI)
bvalmin 0.2            Minimum b-value (up to two) (should be << than expected)
! cvalmin 8.0          minimum column density (log)
! cvalmax 30.0         maximum column density (log)
bltdrop 0.4           drop system if b value less than this AND .. 
clogltdrop 14.3        if log column density less than this. 
bgtdrop 50001.          Drop any system with b .gt. this (> bvalmax, so none).
! colglo 0             column density below which guesses are ignored
fcollallzn 4.162e22    Fik*col for inclusion everywhere.
sigscalemult 1.0       global sigma multiplier
date                   datestamp the '26' summary output
dots 1000              Suppress the dots on screen print.
gcursor                use cursor input for guesses
! ----  add/remove line parameters ----
adsplit 3.0
maxadrem 20
bsigp 10.0            line search lower level for splitting
! adcontf 3.0 10 5.0 12.5
! setspecchar S        internally fixed parameter flag
! zerolevels -0.3 0.3  limits to zero level adjustment (default - none)
! ----  output control
! wr13s f13.9p3w            end summary root filename (fort.13 format)
wr26s f26.9p3w            fort.26 summary root filename
nopchan 1            output (0=not much, 1=iterations to 6, 2=6+fort.18)
vform long                use variable format O/P depending on error estimate
! wrcitn              write current iteration to junk.dat, remove on clean exit
! ----  internal substepping ----
! nsubmin 1           minimum number of subpixels per pixel
nsubmax 3          maximum number of subpixels per data pixel
! nfwhmp 7               require at least 7 pixels per instrument fwhm
! ----  development tools ----
pcvals               prompt for controls (il, stepsize etc)
! verbose
! DEBUG 0
chisqthres 0.001 3.0 0.01      relative chi^2 decrement for stopping
"""

ion_name = dict(o6='OVI 1032', h1='HI 1215', c4='CIV 1548',si4='SiIV 1393',
                ne8='NeVIII 770',n5='NV 1238')
    
atom_format = """\
??    1215.6701 0.416400  6.265E8  1.00 No ID !! MUST BE HERE
>>    1215.6701 0.416400  6.265E8  1.00
<<    1215.6701 0.416400  6.265E8  1.00
<>    1215.6701 0.416400  6.265E8  1.00
__    1215.6701 0.416400  6.265E8  1.00
%s
end   0000.000 0.000000  0.00000   0.000 Terminator
"""

atom_entry = dict(
    o6  = 'O VI  1031.9261 0.13250  4.149e8  M03g',
    h1  = 'H I   1215.6701 0.416400  6.265E8  1.00794 Morton(03)',
    c4  = 'C IV  1548.2041 0.189900 2.642E8 M03 g',
    si4 = 'SiIV  1393.76018 0.513  8.80e8  M03g',
    ne8 = 'NeVIII 770.409 0.103000  1.000e8  VBT',
    n5  = 'N V  1238.821  0.1560  3.391e8   M03g')

########### End info needed for run_vpfit section ##############

# def median_continuum(flux,error,numsig=1.5,plot=False):
#     """ Estimates the continuum using a median and sigma clipping.
    
#     Given the fluxes and one sigma errors for the section, calculates
#     the flux median. Then rejects all flux values less than numsig*sig
#     lower than the median, where sig is the median one sigma error
#     value of the flux values. Repeat this process, only retaining the
#     not-rejected flux values each time, until no flux values are
#     rejected. Once this condition is reached, take the current median
#     value as the continuum.

#     Returns the continuum value.
#     """

#     if plot:  pl.plot(flux)
#     while True:
#         medfl = np.median(flux)
#         meder = np.median(error)
#         #if plot:  l = pl.axhline(medfl)
#         cond = (flux > (medfl - meder*numsig))
#         badflux = flux[~cond]
#         if len(badflux) == 0:
#             return medfl
#         flux = flux[cond]
#         error = error[cond]

def make_vpfit_input(filename, nregions, chsqstop):
    temp = ( 'F ad %(chsqstop)s \n'
             '\n'
             '\n'
             'vpin.dat %(nregions)i\n'
             'n\n'
             '\n' % locals() )
    fh = open(filename,'w')
    fh.write(temp)
    fh.close()

def run_vpfit(specname, wmin, wmax, ion, vpfitexe='vpfit',
              prompts=True, debug=1):

#   create input for running vpfit (fit regions)
    temp = '%%%% %s  1   %10.3f %10.3f\n\n' % (specname,wmin,wmax)
    fh=open('vpin.dat','w')
    fh.writelines(temp)
    fh.close()

    # count number of regions
    f = open('vpin.dat')
    lines = f.readlines()
    f.close()
    nregions = ''.join(lines).count('%%')

    make_vpfit_input('vpfit_input.txt',nregions,chsqstop)

    if debug: print 'vp_input.dat created'

    # Run vpfit
    if prompts:
        raw_input('About to run vpfit - enter to continue')
        
#     temp = ('export ATOMDIR=./atom.dat; export VPFSETUP=./vp_setup.dat;'
#             'echo "atomdir=" $ATOMDIR "vpfsetup=" $VPFSETUP; '
#             '%s < vpfit_input.txt' % (vpfitexe))
    temp = ('%s < vpfit_input.txt' % vpfitexe)
    subprocess.call(temp,shell=1)
    if prompts:
        raw_input('Finished running vpfit - enter to continue')

    # Read in f26 files (one per region)
    # copy all f26 files to one reference file for this spectrum
    f26name = specname[:-4] + '.f26'
    subprocess.call('cat f26* > %s' % f26name, shell=1)
    subprocess.call('rm f26*', shell=1)

def main(ion, nsp, fwhm, s2n, minsig, dv, safe, outprefix, specdir,
         infilename, chsqstop, debug=1):

#   create output directory, and change into that directory
    outdir = os.path.join(specdir,ion)
    if not os.path.lexists(outdir):
        print 'Creating directory %s' % outdir
        os.makedirs(outdir)
        if  debug: print 'Directory created'
    olddir = os.path.abspath('.')
    rundir = olddir
    print 'Changing to dir %s' % outdir
    os.chdir(outdir)

#   make required input files:
    # vp_setup.dat
    fh = open('vp_setup.dat','w')
    fh.writelines(vpsetup_txt + 'guessline %s\n' % ion_name[ion])
    fh.close()

    # atom.dat
    fh = open('atom.dat','w')
    fh.write(atom_format % atom_entry[ion])
    fh.close()

    if debug: print 'atom.dat and vp_setup.dat created'

#   read all spectra for this particular ion, convolve with Gaussian
#   with specified FWHM, rebin to dv (km/s) spectral resolution, add
#   Gaussian noise with signal-to-noise s2n at the continuum extract
#   region of 2*safe (km/s) on either side of spectrum centre
    spectra = astro.utilities.read_specwizardshort(
        infilename, ion=ion, nspec=nsp, mult_tau=1.2,
         s2n=s2n, minsig=minsig, fwhm=fwhm, dv=dv, safe=safe)

    if debug:  print 'finished reading spectra'

    wmin,wmax = spectra.safefit     # fitting region
    wa = spectra.wa
    for s in spectra.spec:
        # write spectra to ascii file
        #contlevel =  median_continuum(s.fl,s.er,numsig=0.9)
        #co = contlevel * np.ones(len(s.fl))
        hd = ['RESVEL %4.2f\n' % fwhm]
        output = ('%15.7f %10.4f %10.4f %10.4f\n' % t
                  for t in zip(wa,s.fl,s.er,s.co))
        outname = '%s%s_%03i.txt' % (outprefix,ion,s.number)
        print 'Writing to %s' % outname
        fh = open(outname,'w')
        fh.writelines(hd + sorted(output))
        fh.close()
        
        if debug: print 'spectrum written'
        
        # fit spectra
        run_vpfit(outname,wmin,wmax,ion,prompts=0,vpfitexe='vpfit')
        
    print 'Changing back to dir %s' % olddir
    os.chdir(olddir)
    return spectra

if 1:
    #spec = main(ion, nsp, fwhm, s2n, minsig, dv, safe, outprefix, specdir,
    #            infilename, chsqstop)
    #outprefix = 'runB_'
    prefix = 'LoRes/Spectra/QSO-correlations4'
    #prefix = 'LoRes/Spectra/QSO-correlations'
    indir = '/media/cosma/data/rw17/nhmc/GIMIC/' + prefix
    suffixes = (#'Sigma0/snapshot_029/Spectrum.Part_los_z1.000.hdf5',
                #'Sigma0/snapshot_039/Spectrum.Part_los_z0.500.hdf5',
                #'Sigma0/snapshot_046/Spectrum.Part_los_z0.250.hdf5', 
                #'Sigmam1/snapshot_029/Spectrum.Part_los_z1.000.hdf5', 
                #'Sigmam1/snapshot_039/Spectrum.Part_los_z0.500.hdf5', 
                #'Sigmam1/snapshot_046/Spectrum.Part_los_z0.250.hdf5', 
                #'Sigmam2/snapshot_027/Spectrum.Part_los_z1.000.hdf5', 
                #'Sigmam2/snapshot_029/Spectrum.Part_los_z0.500.hdf5', 
                #'Sigmam2/snapshot_031/Spectrum.Part_los_z0.250.hdf5', 
                #'Sigmap1/snapshot_029/Spectrum.Part_los_z1.000.hdf5',
                #'Sigmap1/snapshot_039/Spectrum.Part_los_z0.500.hdf5', 
                'Sigmap1/snapshot_046/Spectrum.Part_los_z0.250.hdf5', 
                'Sigmap2/snapshot_033/Spectrum.Part_los_z1.000.hdf5', 
                'Sigmap2/snapshot_043/Spectrum.Part_los_z0.500.hdf5', 
                'Sigmap2/snapshot_051/Spectrum.Part_los_z0.250.hdf5') 

    spectra = []
    for suffix in suffixes:
        specdir = os.path.join(prefix, suffix.split('Spectrum')[0])
        print 'specdir=',specdir
        infilename = os.path.join(indir,suffix)
        spec = main(ion, nsp, fwhm, s2n, minsig, dv, safe, outprefix, specdir,
                    infilename, chsqstop)
        spectra.append(spec)
